#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
from bs4 import BeautifulSoup
import sys
from tqdm import tqdm


def DescriptorAnalyzer(file):
    DB_NAME = "mesh"
    COLLECTION_NAME = file.split(".")[0]

    connection = MongoClient('localhost', 27017)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]

    infile = open(file, "r")
    currentString = ""
    for line in tqdm(infile):
        line = line.rstrip()
        currentString += line
        if line.find("</DescriptorRecord>") != -1:
            record = {}
            soup = BeautifulSoup(currentString, "xml")
            did = soup.DescriptorUI
            if did != None:
                record["DescriptorID"] = did.string
            if soup.find("AllowableQualifiersList") != None:
                aqlist = soup.AllowableQualifiersList.find_all("QualifierUI")
                if aqlist != None:
                    allowQual = []
                    for aq in aqlist:
                        allowQual.append(aq.string)
                    record["AllowableQualifiersList"] = allowQual
            if soup.find("TreeNumberList") != None:
                tnlist = soup.TreeNumberList.find_all("TreeNumber")
                if tnlist != None:
                    treeNumber = []
                    for tn in tnlist:
                        treeNumber.append(tn.string)
                    record["TreeNumberList"] = treeNumber
            anno = soup.find("Annotation")
            if anno != None:
                record["Annotation"] = anno.string
            if soup.find("ConceptList") != None:
                tlist = soup.ConceptList.find_all("Term")
                if tlist != None:
                    termList = []
                    for t in tlist:
                        if t["IsPermutedTermYN"] == "N":
                            termList.append(t.String.string)
                            if t["RecordPreferredTermYN"] == "Y":
                                record["DescriptorName"] = t.String.string
                    record["TermList"] = termList
            collection.insert_one(record)
            currentString = ""
    infile.close()
    connection.close()


def QualifierAnalyzer(file):
    DB_NAME = "mesh"
    COLLECTION_NAME = file.split(".")[0]

    connection = MongoClient('localhost', 27017)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]

    infile = open(file, "r")
    currentString = ""
    for line in tqdm(infile):
        line = line.rstrip()
        currentString += line
        if line.find("</QualifierRecord>") != -1:
            record = {}
            soup = BeautifulSoup(currentString, "xml")
            qid = soup.QualifierUI
            if qid != None:
                record["QualifierID"] = qid.string
            if soup.find("TreeNumberList") != None:
                tnlist = soup.TreeNumberList.find_all("TreeNumber")
                if tnlist != None:
                    treeNumber = []
                    for tn in tnlist:
                        treeNumber.append(tn.string)
                    record["TreeNumberList"] = treeNumber
            anno = soup.find("Annotation")
            if anno != None:
                record["Annotation"] = anno.string
            if soup.find("ConceptList") != None:
                tlist = soup.ConceptList.find_all("Term")
                if tlist != None:
                    termList = []
                    for t in tlist:
                        if t["IsPermutedTermYN"] == "N":
                            termList.append(t.String.string)
                            if t["RecordPreferredTermYN"] == "Y":
                                record["QualifierName"] = t.String.string
                    record["TermList"] = termList
            collection.insert_one(record)
            currentString = ""
    infile.close()
    connection.close()


def SupplementalAnalyzer(file):
    DB_NAME = "mesh"
    COLLECTION_NAME = file.split(".")[0]

    connection = MongoClient('localhost', 27017)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]

    infile = open(file, "r")
    currentString = ""
    for line in tqdm(infile):
        line = line.rstrip()
        currentString += line
        if line.find("</SupplementalRecord>") != -1:
            record = {}
            soup = BeautifulSoup(currentString, "xml")
            sid = soup.QualifierUI
            if sid != None:
                record["SupplementalRecordID"] = sid.string
            note = soup.Note
            if note != None:
                record["Note"] = note.string
            if soup.find("ConceptList") != None:
                tlist = soup.ConceptList.find_all("Term")
                if tlist != None:
                    termList = []
                    for t in tlist:
                        if t["IsPermutedTermYN"] == "N":
                            termList.append(t.String.string)
                            if t["RecordPreferredTermYN"] == "Y":
                                record["SupplementalRecordName"] = t.String.string
                    record["TermList"] = termList
            collection.insert_one(record)
            currentString = ""
    infile.close()
    connection.close()


if __name__ == "__main__":
    year = sys.argv[1]
    DescriptorAnalyzer("desc" + year + ".xml")
    QualifierAnalyzer("qual" + year + ".xml")
    SupplementalAnalyzer("supp" + year + ".xml")
